# -*- coding: utf-8 -*-
from elasticsearch import Elasticsearch
from wt.connector.data_connect import DataConnect, ElasticSearchOperate
from wt.connector.db_manager.db_untils_tools import DBUntilsTools

class DBElasticSearchManager(DataConnect, DBUntilsTools):
    def __init__(self):
        super().__init__()

    def generate_create_ddl(self, db, index_name):
        """
        定义 Elasticsearch 索引映射
        :return: 索引映射
        """
        mappings = {
            "mappings": {
                "properties": {
                    "id": {"type": "integer"},
                    "c_02_integer_col": {"type": "integer"},
                    "c_03_double_col": {"type": "double"},
                    "c_04_float_col": {"type": "float"},
                    "c_05_smallint_col": {"type": "short"},
                    "c_06_bigint_col": {"type": "long"},
                    "c_07_tinyint_col": {"type": "byte"},
                    "c_08_decimal_col": {"type": "scaled_float", "scaling_factor": 100},
                    "c_09_varchar_col": {"type": "keyword"},
                    "c_10_date_col": {"type": "date"},
                    "c_11_datetime_col": {"type": "date"}
                }
            }
        }
        return ElasticSearchOperate(hasindices=True,table=index_name,operate="create", data=mappings)

    def insert_document(self, index_name, num):
        """
        插入文档到 Elasticsearch 索引
        :param index_name: 索引名称
        :param doc: 文档内容
        :return:
        """
        for _ in range(num):
            doc = self.example_document()
            res = self.exec_es_sql(None, ElasticSearchOperate(table=index_name,operate="insert", data=doc))
            print(f"Document inserted: {res}")

    def example_document(self):
        """
        返回一个示例文档
        :return: 示例文档
        """
        return {
            "id": self.get_random_id(),
            "c_02_integer_col": self.get_random(-2147483648, 2147483648),
            "c_03_double_col": self.get_random_float(-500, 500),
            "c_04_float_col": self.get_random_float(-500, 500),
            "c_05_smallint_col": self.get_random(-32768, 32767),
            "c_06_bigint_col": self.get_random(-9223372036854775808, 9223372036854775807),
            "c_07_tinyint_col": self.get_random(-128, 123),
            "c_08_decimal_col": self.get_ramdom_decimal(2, 2),
            "c_09_varchar_col": self.get_random_string(8),
            "c_10_date_col": self.get_random_date(),
            "c_11_datetime_col": self.get_random_datetime()
        }

    def show_index(self, index_name):
        """
        显示 Elasticsearch 索引的映射
        :param index_name: 索引名称
        :return:
        """
        res = self.exec_es_sql(None, ElasticSearchOperate(hasindices=True,table=index_name,operate="exists", data=None))
        if res['message']:
            mapping = self.exec_es_sql(None, ElasticSearchOperate(hasindices=True,table=index_name,operate="get_mapping", data=None))
            print(f"Mapping for index {index_name}: {mapping}")
        else:
            print(f"Index {index_name} does not exist")

    def execute_drop_table(self, db, index_name):
        """
        删除 Elasticsearch 索引
        :param index_name: 索引名称
        :return:
        """
        print("execute_drop_table")
        res = self.exec_es_sql(None, ElasticSearchOperate(hasindices=True,table=index_name,operate="exists", data=None))
        if res['message']:
            res = self.exec_es_sql(None, ElasticSearchOperate(hasindices=True,table=index_name,operate="delete", data=None))
            if res['data'] == "True":
                return res['data']
            else:
                return None
        else:
            print(f"Index {index_name} does not exist")

    def get_count(self, db, index_name):
        """
        获取 Elasticsearch 索引下的文档数量
        :param index_name: 索引名称
        :return: 文档数量
        """
        try:
            res = self.exec_es_sql(None, ElasticSearchOperate(table=index_name,operate="count", data=None))
            count = res['message']['count']
            return count
        except Exception as e:
            print(f"Failed to get document count from index {index_name}: {str(e)}")
            return None

if __name__ == '__main__':
    es_manager = DBElasticSearchManager()
    es_manager.create_and_insert_sql("auto_test_20240703",10)
